import argparse
import json
import os
import sys
import time
from pathlib import Path
from collections import defaultdict
from typing import Dict, List, Tuple

from executor import SpanLinkExecutor


class ExperimentRunner:
    """实验运行器"""

    def __init__(self, data_file: str, verbose: bool = False):
        """
        初始化实验运行器
        
        Args:
            data_file: 数据文件路径
            verbose: 是否显示详细信息
        """
        self.data_file = data_file
        self.verbose = verbose
        
        # 实验参数
        self.weight_methods = ['idf', 'entropy', 'bm25']
        self.weight_topks = [1, 3, 5]
        
        # 实验结果存储
        self.exp_results = defaultdict(dict)  # exp_results[weight_method][weight_topk] = acc
        self.detailed_results = defaultdict(dict)  # 详细结果，包含更多指标
        
    def run_single_experiment(self, weight_method: str, weight_topk: int) -> Dict:
        """
        运行单个实验
        
        Args:
            weight_method: 权重计算方法
            weight_topk: Top-K 参数
            
        Returns:
            Dict: 实验结果，包含准确率等指标
        """
        print(f"\n{'='*80}")
        print(f"运行实验: weight_method={weight_method}, weight_topk={weight_topk}")
        print(f"{'='*80}")
        
        start_time = time.time()
        
        try:
            # 创建执行器
            executor = SpanLinkExecutor(data_file=self.data_file, blocked_services=[], weight_method=weight_method,
                                        weight_topk=weight_topk)
            
            # 运行流水线
            executor.run_pipeline()
            
            # 提取结果
            result = self._extract_results(executor)
            result['execution_time'] = time.time() - start_time
            result['weight_method'] = weight_method
            result['weight_topk'] = weight_topk
            
            print(f"\n实验完成! 准确率: {result['accuracy']:.2%}, 耗时: {result['execution_time']:.2f}秒")
            
            return result
            
        except Exception as e:
            print(f"实验失败: {e}")
            import traceback
            traceback.print_exc()
            
            return {
                'accuracy': 0.0,
                'total_predictions': 0,
                'correct_predictions': 0,
                'ground_truth': 0,
                'execution_time': time.time() - start_time,
                'weight_method': weight_method,
                'weight_topk': weight_topk,
                'error': str(e)
            }
    
    def _extract_results(self, executor: SpanLinkExecutor) -> Dict:
        """
        从执行器中提取实验结果
        
        Args:
            executor: 执行器实例
            
        Returns:
            Dict: 实验结果
        """
        # 计算准确率（复制 executor.py 中的评估逻辑）
        checkout_predictions = {}
        checkout_ground_truth = {}
        
        # 收集 checkout 相关的预测和真实关联
        for out_ep, assignments in executor.pred_assignments.items():
            if "checkout" in out_ep.lower():
                checkout_predictions[out_ep] = assignments
        
        for out_ep, assignments in executor.true_assignments.items():
            if "checkout" in out_ep.lower():
                checkout_ground_truth[out_ep] = assignments
        
        # 计算指标
        ground_truth = sum(len(assignments) for assignments in checkout_ground_truth.values())
        total_predictions = sum(len(assignments) for assignments in checkout_predictions.values())
        
        correct_predictions = 0
        for out_ep in checkout_predictions:
            pred_assignments = checkout_predictions[out_ep]
            true_assignments = checkout_ground_truth.get(out_ep, {})
            
            for in_span_id, out_span_id in pred_assignments.items():
                if in_span_id in true_assignments:
                    true_out_span_id = true_assignments[in_span_id]
                    if out_span_id == true_out_span_id:
                        correct_predictions += 1
        
        accuracy = correct_predictions / ground_truth if ground_truth > 0 else 0.0
        
        return {
            'accuracy': accuracy,
            'total_predictions': total_predictions,
            'correct_predictions': correct_predictions,
            'ground_truth': ground_truth,
            'precision': correct_predictions / total_predictions if total_predictions > 0 else 0.0,
            'recall': correct_predictions / ground_truth if ground_truth > 0 else 0.0
        }
    
    def run_all_experiments(self):
        """运行所有实验组合"""
        print("开始运行对比实验...")
        print(f"权重方法: {self.weight_methods}")
        print(f"Top-K 参数: {self.weight_topks}")
        print(f"总实验数: {len(self.weight_methods) * len(self.weight_topks)}")
        
        total_start_time = time.time()
        experiment_count = 0
        
        for weight_method in self.weight_methods:
            for weight_topk in self.weight_topks:
                experiment_count += 1
                print(f"\n进度: {experiment_count}/{len(self.weight_methods) * len(self.weight_topks)}")
                
                # 运行单个实验
                result = self.run_single_experiment(weight_method, weight_topk)
                
                # 存储结果
                self.exp_results[weight_method][weight_topk] = result['accuracy']
                self.detailed_results[weight_method][weight_topk] = result
        
        total_time = time.time() - total_start_time
        print(f"\n所有实验完成! 总耗时: {total_time:.2f}秒")
    
    def print_results_table(self):
        """打印结果表格"""
        print(f"\n{'='*80}")
        print("实验结果汇总 - 准确率 (Accuracy)")
        print(f"{'='*80}")
        
        # 表头
        header = f"{'Weight Method':<15}"
        for topk in self.weight_topks:
            header += f"{'Top-' + str(topk):<12}"
        print(header)
        print("-" * len(header))
        
        # 数据行
        for method in self.weight_methods:
            row = f"{method:<15}"
            for topk in self.weight_topks:
                acc = self.exp_results[method][topk]
                row += f"{acc:<12.2%}"
            print(row)
        
        # 找出最佳组合
        best_acc = 0.0
        best_method = None
        best_topk = None
        
        for method in self.weight_methods:
            for topk in self.weight_topks:
                acc = self.exp_results[method][topk]
                if acc > best_acc:
                    best_acc = acc
                    best_method = method
                    best_topk = topk
        
        print(f"\n最佳组合: weight_method={best_method}, weight_topk={best_topk}, accuracy={best_acc:.2%}")
    
    def print_detailed_results(self):
        """打印详细结果"""
        print(f"\n{'='*80}")
        print("详细实验结果")
        print(f"{'='*80}")
        
        for method in self.weight_methods:
            print(f"\n权重方法: {method.upper()}")
            print("-" * 60)
            
            for topk in self.weight_topks:
                result = self.detailed_results[method][topk]
                print(f"  Top-K={topk}:")
                print(f"    准确率 (Accuracy): {result['accuracy']:.2%}")
                print(f"    精确率 (Precision): {result['precision']:.2%}")
                print(f"    召回率 (Recall): {result['recall']:.2%}")
                print(f"    总预测数: {result['total_predictions']}")
                print(f"    正确预测数: {result['correct_predictions']}")
                print(f"    真值数: {result['ground_truth']}")
                print(f"    执行时间: {result['execution_time']:.2f}秒")
                if 'error' in result:
                    print(f"    错误: {result['error']}")
                print()
    
    def save_results(self, output_file: str):
        """
        保存实验结果到文件
        
        Args:
            output_file: 输出文件路径
        """
        # 准备保存的数据
        save_data = {
            'experiment_config': {
                'data_file': self.data_file,
                'weight_methods': self.weight_methods,
                'weight_topks': self.weight_topks
            },
            'accuracy_results': dict(self.exp_results),
            'detailed_results': dict(self.detailed_results),
            'summary': {
                'total_experiments': len(self.weight_methods) * len(self.weight_topks),
                'best_combination': self._find_best_combination()
            }
        }
        
        # 创建输出目录
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 保存到 JSON 文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(save_data, f, ensure_ascii=False, indent=2)
        
        print(f"\n实验结果已保存到: {output_file}")
    
    def _find_best_combination(self) -> Dict:
        """找出最佳参数组合"""
        best_acc = 0.0
        best_method = None
        best_topk = None
        
        for method in self.weight_methods:
            for topk in self.weight_topks:
                acc = self.exp_results[method][topk]
                if acc > best_acc:
                    best_acc = acc
                    best_method = method
                    best_topk = topk
        
        return {
            'weight_method': best_method,
            'weight_topk': best_topk,
            'accuracy': best_acc
        }


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='运行权重方法和 Top-K 参数的对比实验')
    parser.add_argument('--data_file', type=str,
                        default='data/cbt/checkout_100_user50.json',
                        help='数据文件路径')
    parser.add_argument('--output_file', type=str,
                        default='data/experiment_results.json',
                        help='实验结果输出文件路径')
    parser.add_argument('--verbose', action='store_true',
                        help='显示详细信息')
    parser.add_argument('--show_detailed', action='store_true',
                        help='显示详细实验结果')
    
    args = parser.parse_args()
    
    # 检查数据文件是否存在
    if not os.path.exists(args.data_file):
        print(f"错误: 数据文件 {args.data_file} 不存在")
        sys.exit(1)
    
    # 创建实验运行器
    runner = ExperimentRunner(args.data_file, args.verbose)
    
    try:
        # 运行所有实验
        runner.run_all_experiments()
        
        # 打印结果
        runner.print_results_table()
        
        if args.show_detailed:
            runner.print_detailed_results()
        
        # 保存结果
        runner.save_results(args.output_file)
        
        print(f"\n{'='*80}")
        print("对比实验完成!")
        print(f"{'='*80}")
        
    except KeyboardInterrupt:
        print("\n实验被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"实验运行出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()